package org.wzq.gradle_learn;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.spec.SecretKeySpec;
import java.io.*;
import java.security.InvalidKeyException;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;


/**
 * AES加密解密类,注这个类是与服务端搭配使用的 nodejs加密，这里解密，与JAVA加解密有一些区别
 * @author wzq
 */
public class AesTool {
    private static final String CHARSET_UTF8 = "UTF-8";
    private Cipher encryptCipher = null;// 加密
    private Cipher decryptCipher = null;// 解密

    public AesTool(String key) {
        try {
            byte[] raw = getByteKey(key);
            init(raw);
        } catch(Exception e) {
            e.printStackTrace();
        }
    }


    private byte[] getByteKey(String key) {
        byte[] raw = null;
        try {
            byte[] keyb = key.getBytes(CHARSET_UTF8);
            MessageDigest md = MessageDigest.getInstance("MD5");
            raw = md.digest(keyb);
        } catch(NoSuchAlgorithmException e) {
            //e.printStackTrace();
        } catch(UnsupportedEncodingException e) {
            //e.printStackTrace();
        }
        return raw;
    }

    private void init(byte[] key) {
        try {
            SecretKeySpec skeySpec = new SecretKeySpec(key, "AES");
            //
            encryptCipher = Cipher.getInstance("AES");
            encryptCipher.init(Cipher.ENCRYPT_MODE, skeySpec);
            //
            decryptCipher = Cipher.getInstance("AES");
            decryptCipher.init(Cipher.DECRYPT_MODE, skeySpec);
        } catch(NoSuchAlgorithmException e) {
            //e.printStackTrace();
        } catch(InvalidKeyException e) {
            //e.printStackTrace();
        } catch(NoSuchPaddingException e) {
            //e.printStackTrace();
        }
    }

    /**
     * 加密byte[]
     * @param bytes
     * @return
     */
    private byte[] encryptBytes(byte[] bytes) {
        try {
            return encryptCipher.doFinal(bytes);
        } catch(IllegalBlockSizeException e) {
            //e.printStackTrace();
        } catch(BadPaddingException e) {
            //e.printStackTrace();
        }
        return null;
    }

    /**
     * 解密byte[]
     * @param src
     * @return
     */
    private byte[] decryptBytes(byte[] src) {
        try {
            return decryptCipher.doFinal(src);
        } catch(IllegalBlockSizeException e) {
            //e.printStackTrace();
        } catch(BadPaddingException e) {
            //e.printStackTrace();
        }
        return null;
    }


    private InputStream encryptStream(InputStream in) {
        try {
            // InputStream => ByteArray => encrypt bytes => InputStream
            byte[] bytes = stream2Bytes(in);
            byte[] encryptBytes = encryptBytes(bytes);
            return bytes2Stream(encryptBytes);
        } catch(IOException e) {
            //e.printStackTrace();
            return null;
        }
    }

    private InputStream decryptStream(InputStream in) {
        try {
            // InputStream => ByteArray => decrypt bytes => InputStream
            byte[] bytes = stream2Bytes(in);
            byte[] decryptBytes = decryptBytes(bytes);
            return bytes2Stream(decryptBytes);
        } catch(IOException e) {
            //e.printStackTrace();
            return null;
        }
    }

    // 对文件加密，替换原文件
    public void encryptFile(String fileName) throws IOException {
        File f = new File(fileName);
        InputStream src = new FileInputStream(f);
        InputStream dst = encryptStream(src);
        f.delete();

        File fileSave = new File(fileName);
        FileOutputStream fos = new FileOutputStream(fileSave);
        fos.write(stream2Bytes(dst));
        src.close();
        dst.close();
        fos.close();
    }

    // 对src文件加密，存放dst
    public void encryptFile(String src, String dst) throws IOException {
        File srcFile = new File(src);
        InputStream srcStream = new FileInputStream(srcFile);
        InputStream dstStream = encryptStream(srcStream);

        // 建立目录
        File dstFile = new File(dst);
        if(!dstFile.getParentFile().exists()) {
            dstFile.getParentFile().mkdirs();
        }

        FileOutputStream fos = new FileOutputStream(dstFile);
        fos.write(stream2Bytes(dstStream));
        fos.flush();
        fos.close();
    }

    public void decryptFile(String src, String dst) throws IOException {
        File srcFile = new File(src);
        InputStream srcStream = new FileInputStream(srcFile);
        InputStream dstStream = decryptStream(srcStream);

        // 建立目录
        File dstFile = new File(dst);
        if(!dstFile.getParentFile().exists()) {
            dstFile.getParentFile().mkdirs();
        }

        FileOutputStream fos = new FileOutputStream(dstFile);
        fos.write(stream2Bytes(dstStream));
        fos.flush();
        fos.close();
    }

    /**
     * inputStream => byte[]
     * @param is
     * @return
     * @throws java.io.IOException
     */
    public static byte[] stream2Bytes(InputStream is) throws IOException {
        ByteArrayOutputStream buffer = new ByteArrayOutputStream();
        int nRead;
        final int BUFFER_SIZE = 1024;
        byte[] data = new byte[BUFFER_SIZE];
        while((nRead = is.read(data, 0, data.length)) != -1) {
            buffer.write(data, 0, nRead);
        }
        buffer.flush();
        byte[] bytes = buffer.toByteArray();
        return bytes;
    }

    /**
     * byte[] => InputStream
     * @param in
     * @throws Exception
     */
    public static InputStream bytes2Stream(byte[] in) throws IOException {
        if(in == null) {
            return null;
        }
        ByteArrayInputStream is = new ByteArrayInputStream(in);
        return is;
    }
}
